// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/chunked_array.h>

#include <algorithm>
#include <cstdlib>
#include <memory>
#include <sstream>
#include <utility>

#include <nebula/array/array_base.h>
#include <nebula/array/array_nested.h>
#include <nebula/array/util.h>
#include <nebula/array/validate.h>
#include <nebula/core/pretty_print.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/checked_cast.h>
#include <turbo/log/logging.h>

namespace nebula {

    class MemoryPool;

    // ----------------------------------------------------------------------
    // ChunkedArray methods

    ChunkedArray::ChunkedArray(ArrayVector chunks, std::shared_ptr<DataType> type)
            : chunks_(std::move(chunks)),
              type_(std::move(type)),
              length_(0),
              null_count_(0),
              chunk_resolver_{chunks_} {
        if (type_ == nullptr) {
            KCHECK_GT(chunks_.size(), 0)
            << "cannot construct ChunkedArray from empty vector and omitted type";
            type_ = chunks_[0]->type();
        }

        for (const auto &chunk: chunks_) {
            length_ += chunk->length();
            null_count_ += chunk->null_count();
        }
    }

    turbo::Result<std::shared_ptr<ChunkedArray>> ChunkedArray::create(ArrayVector chunks,
                                                                    std::shared_ptr<DataType> type) {
        if (type == nullptr) {
            if (chunks.size() == 0) {
                return turbo::invalid_argument_error(
                        "cannot construct ChunkedArray from empty vector "
                        "and omitted type");
            }
            type = chunks[0]->type();
        }
        for (const auto &chunk: chunks) {
            if (!chunk->type()->equals(*type)) {
                return turbo::failed_precondition_error("Array chunks must all be same type");
            }
        }
        return std::make_shared<ChunkedArray>(std::move(chunks), std::move(type));
    }

    turbo::Result<std::shared_ptr<ChunkedArray>> ChunkedArray::make_empty(
            std::shared_ptr<DataType> type, MemoryPool *memory_pool) {
        std::vector<std::shared_ptr<Array>> new_chunks(1);
        TURBO_MOVE_OR_RAISE(new_chunks[0], MakeEmptyArray(type, memory_pool));
        return std::make_shared<ChunkedArray>(std::move(new_chunks));
    }

    bool ChunkedArray::equals(const ChunkedArray &other, const EqualOptions &opts) const {
        if (length_ != other.length()) {
            return false;
        }
        if (null_count_ != other.null_count()) {
            return false;
        }
        // We cannot toggle check_metadata here yet, so we don't check it
        if (!type_->equals(*other.type_, /*check_metadata=*/false)) {
            return false;
        }

        // Check contents of the underlying arrays. This checks for equality of
        // the underlying data independently of the chunk size.
        return internal::ApplyBinaryChunked(
                *this, other,
                [&](const Array &left_piece, const Array &right_piece,
                    int64_t TURBO_ARG_UNUSED(position)) {
                    if (!left_piece.equals(right_piece, opts)) {
                        return turbo::invalid_argument_error("Unequal piece");
                    }
                    return turbo::OkStatus();
                })
                .ok();
    }

    namespace {

        bool mayHaveNaN(const nebula::DataType &type) {
            if (type.num_fields() == 0) {
                return is_floating(type.id());
            } else {
                for (const auto &field: type.fields()) {
                    if (mayHaveNaN(*field->type())) {
                        return true;
                    }
                }
            }
            return false;
        }

    }  //  namespace

    bool ChunkedArray::equals(const std::shared_ptr<ChunkedArray> &other,
                              const EqualOptions &opts) const {
        if (!other) {
            return false;
        }
        if (this == other.get() && !mayHaveNaN(*type_)) {
            return true;
        }
        return equals(*other.get(), opts);
    }

    bool ChunkedArray::approx_equals(const ChunkedArray &other,
                                    const EqualOptions &equal_options) const {
        if (length_ != other.length()) {
            return false;
        }
        if (null_count_ != other.null_count()) {
            return false;
        }
        // We cannot toggle check_metadata here yet, so we don't check it
        if (!type_->equals(*other.type_, /*check_metadata=*/false)) {
            return false;
        }

        // Check contents of the underlying arrays. This checks for equality of
        // the underlying data independently of the chunk size.
        return internal::ApplyBinaryChunked(
                *this, other,
                [&](const Array &left_piece, const Array &right_piece,
                    int64_t TURBO_ARG_UNUSED(position)) {
                    if (!left_piece.approx_equals(right_piece, equal_options)) {
                        return turbo::invalid_argument_error("Unequal piece");
                    }
                    return turbo::OkStatus();
                })
                .ok();
    }

    turbo::Result<std::shared_ptr<Scalar>> ChunkedArray::get_scalar(int64_t index) const {
        const auto loc = chunk_resolver_.resolve(index);
        if (loc.chunk_index >= static_cast<int64_t>(chunks_.size())) {
            return turbo::out_of_range_error("index with value of ", index,
                                             " is out-of-bounds for chunked array of length ", length_);
        }
        return chunks_[loc.chunk_index]->get_scalar(loc.index_in_chunk);
    }

    std::shared_ptr<ChunkedArray> ChunkedArray::slice(int64_t offset, int64_t length) const {
        KCHECK_LE(offset, length_) << "slice offset greater than array length";
        bool offset_equals_length = offset == length_;
        int curr_chunk = 0;
        while (curr_chunk < num_chunks() && offset >= chunk(curr_chunk)->length()) {
            offset -= chunk(curr_chunk)->length();
            curr_chunk++;
        }

        ArrayVector new_chunks;
        if (num_chunks() > 0 && (offset_equals_length || length == 0)) {
            // Special case the zero-length slice to make sure there is at least 1 Array
            // in the result. When there are zero chunks we return zero chunks
            new_chunks.push_back(chunk(std::min(curr_chunk, num_chunks() - 1))->slice(0, 0));
        } else {
            while (curr_chunk < num_chunks() && length > 0) {
                new_chunks.push_back(chunk(curr_chunk)->slice(offset, length));
                length -= chunk(curr_chunk)->length() - offset;
                offset = 0;
                curr_chunk++;
            }
        }

        return std::make_shared<ChunkedArray>(new_chunks, type_);
    }

    std::shared_ptr<ChunkedArray> ChunkedArray::slice(int64_t offset) const {
        return slice(offset, length_);
    }

    turbo::Result<std::vector<std::shared_ptr<ChunkedArray>>> ChunkedArray::Flatten(
            MemoryPool *pool) const {
        if (type()->id() != Type::STRUCT) {
            // Emulate nonexistent copy constructor
            return std::vector<std::shared_ptr<ChunkedArray>>{
                    std::make_shared<ChunkedArray>(chunks_, type_)};
        }

        std::vector<ArrayVector> flattened_chunks(type()->num_fields());
        for (const auto &chunk: chunks_) {
            TURBO_MOVE_OR_RAISE(auto arrays,
                                turbo::checked_cast<const StructArray &>(*chunk).Flatten(pool));
                    DKCHECK_EQ(arrays.size(), flattened_chunks.size());
            for (size_t i = 0; i < arrays.size(); ++i) {
                flattened_chunks[i].push_back(arrays[i]);
            }
        }

        std::vector<std::shared_ptr<ChunkedArray>> flattened(type()->num_fields());
        for (size_t i = 0; i < flattened.size(); ++i) {
            auto child_type = type()->field(static_cast<int>(i))->type();
            flattened[i] =
                    std::make_shared<ChunkedArray>(std::move(flattened_chunks[i]), child_type);
        }
        return flattened;
    }

    turbo::Result<std::shared_ptr<ChunkedArray>> ChunkedArray::View(
            const std::shared_ptr<DataType> &type) const {
        ArrayVector out_chunks(this->num_chunks());
        for (int i = 0; i < this->num_chunks(); ++i) {
            TURBO_MOVE_OR_RAISE(out_chunks[i], chunks_[i]->View(type));
        }
        return std::make_shared<ChunkedArray>(out_chunks, type);
    }

    std::string ChunkedArray::to_string() const {
        std::stringstream ss;
        KCHECK_OK(pretty_print(*this, 0, &ss));
        return ss.str();
    }

    namespace {

        turbo::Status ValidateChunks(const ArrayVector &chunks, bool full_validation) {
            if (chunks.size() == 0) {
                return turbo::OkStatus();
            }

            const auto &type = *chunks[0]->type();
            // Make sure chunks all have the same type
            for (size_t i = 1; i < chunks.size(); ++i) {
                const Array &chunk = *chunks[i];
                if (!chunk.type()->equals(type)) {
                    return turbo::invalid_argument_error("In chunk ", i, " expected type ", type.to_string(),
                                                  " but saw ", chunk.type()->to_string());
                }
            }
            // Validate the chunks themselves
            for (size_t i = 0; i < chunks.size(); ++i) {
                const Array &chunk = *chunks[i];
                const turbo::Status st = full_validation ? internal::ValidateArrayFull(chunk)
                                                         : internal::ValidateArray(chunk);
                if (!st.ok()) {
                    return turbo::invalid_argument_error("In chunk ", i, ": ", st.to_string());
                }
            }
            return turbo::OkStatus();
        }

    }  // namespace

    turbo::Status ChunkedArray::validate() const {
        return ValidateChunks(chunks_, /*full_validation=*/false);
    }

    turbo::Status ChunkedArray::validate_full() const {
        return ValidateChunks(chunks_, /*full_validation=*/true);
    }

    namespace internal {

        bool MultipleChunkIterator::next(std::shared_ptr<Array> *next_left,
                                         std::shared_ptr<Array> *next_right) {
            if (pos_ == length_) return false;

            // Find non-empty chunk
            std::shared_ptr<Array> chunk_left, chunk_right;
            while (true) {
                chunk_left = left_.chunk(chunk_idx_left_);
                chunk_right = right_.chunk(chunk_idx_right_);
                if (chunk_pos_left_ == chunk_left->length()) {
                    chunk_pos_left_ = 0;
                    ++chunk_idx_left_;
                    continue;
                }
                if (chunk_pos_right_ == chunk_right->length()) {
                    chunk_pos_right_ = 0;
                    ++chunk_idx_right_;
                    continue;
                }
                break;
            }
            // Determine how big of a section to return
            int64_t iteration_size = std::min(chunk_left->length() - chunk_pos_left_,
                                              chunk_right->length() - chunk_pos_right_);

            *next_left = chunk_left->slice(chunk_pos_left_, iteration_size);
            *next_right = chunk_right->slice(chunk_pos_right_, iteration_size);

            pos_ += iteration_size;
            chunk_pos_left_ += iteration_size;
            chunk_pos_right_ += iteration_size;
            return true;
        }

    }  // namespace internal
}  // namespace nebula
